package com.ruleengine.worker.runner.chain;

import java.util.Comparator;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import com.ruleengine.destination.DestinationManager;
import com.ruleengine.destination.IDestinationProvider;
import com.ruleengine.worker.model.Chain;
import com.ruleengine.util.model.Metadata;
import com.ruleengine.worker.model.Relation;
import com.ruleengine.worker.model.Rule;
import com.ruleengine.worker.runner.rule.PubSubRuleRunner;
import com.ruleengine.worker.runner.rule.RuleRunner;
import com.ruleengine.worker.runner.rule.RuleRunnerFactory;

import lombok.extern.slf4j.Slf4j;

/**
 * @author lianghanwei on 2024/3/3
 */
@Slf4j
public class ChainRunnerUtil {

    private static final int DEFAULT_MAX_LINKED_RUNNERS_CLOSE_WAIT_TIME_IN_SECONDS = 10;

    public static Rule createDispatchRule(String name) {
        Metadata metadata = Metadata.builder()
            .name(name)
            .build();

        Rule.Spec spec = Rule.Spec.builder()
            .type(Rule.RuleType.DISPATCH)
            .build();

        return Rule.builder()
            .spec(spec)
            .metadata(metadata)
            .build();
    }


    public static void initChainRunner(Map<String, RuleRunner> ruleUidRunnerMap, Chain chain,
                                       PubSubRuleRunner dispatchRunner, DestinationManager destinationManager) {
        ruleUidRunnerMap.clear();
        createRuleRunner(ruleUidRunnerMap, chain.getSpec().getRules(), destinationManager);
        linkRuleRunners(ruleUidRunnerMap, chain.getSpec().getRelations(), dispatchRunner);
    }

    private static void createRuleRunner(Map<String,
                                         RuleRunner> ruleUidRunnerMap,
                                         List<Rule> rules,
                                         DestinationManager destinationManager) {
        for (Rule rule : rules) {
            RuleRunner ruleRunner = RuleRunnerFactory.newInstance(rule, destinationManager);
            ruleUidRunnerMap.putIfAbsent(rule.getMetadata().getUid(), ruleRunner);
        }
    }

    private static void linkRuleRunners(Map<String, RuleRunner> ruleUidRunnerMap,
                                        List<Relation> relations,
                                        RuleRunner dispatchRunner) {
        Set<String> pubRunnerUidSet = new HashSet<>();
        Set<String> subRunnerUidSet = new HashSet<>();

        for (Relation relation : relations) {
            String upperUid = relation.getSpec().getUpper().getMetadata().getUid();
            String downUid = relation.getSpec().getDown().getMetadata().getUid();

            // TODO multi link between each runner is not supported
            // build link between each runner
            if (ruleUidRunnerMap.containsKey(upperUid) && ruleUidRunnerMap.containsKey(downUid)) {
                ruleUidRunnerMap.get(upperUid).subscribe(ruleUidRunnerMap.get(downUid));

                pubRunnerUidSet.add(upperUid);
                subRunnerUidSet.add(downUid);
            }
        }
        linkDispatchRunner(ruleUidRunnerMap, pubRunnerUidSet, subRunnerUidSet, dispatchRunner);
    }

    private static void linkDispatchRunner(Map<String, RuleRunner> ruleUidRunnerMap,
                                           Set<String> pubRunnerUidSet,
                                           Set<String> subRunnerUidSet,
                                           RuleRunner dispatcher) {
        // TODO flag the first rule runner
        // TODO is dispatcher is necessary, it could be the first rule runner
        for (String ruleUid : ruleUidRunnerMap.keySet()) {
            // connect the first rule runner with the chain by `onSubscribe` with the chain runner
            // the chain runner is the publisher which will submit items to those subscribed rule runners
            // e.g.
            //                 item                item
            //   [DispatchRunner] ----> [RuleRunner1] ----> [RuleRunner2]
            //
            //   No process rule runner(the rule relation is empty), dispatcher rule runner dispatch item directly
            //   [DispatchRunner] ----> [RuleRunner1]
            if (!subRunnerUidSet.contains(ruleUid)) {
                dispatcher.subscribe(ruleUidRunnerMap.get(ruleUid));
            }
        }
    }

    /**
     * Close linked runners
     *
     * Attention: the executors applied during async close should be diff with rule runner executors,
     * otherwise it may block rule runners to consume the left messages
     */
    public static CompletableFuture<Void> closeLinkedRunners(String chainName,
        List<PubSubRuleRunner> ruleRunners, PubSubRuleRunner dispatchRunner, ExecutorService executor) {
        return CompletableFuture.runAsync(() -> {
                int lag = getMaxLinkedRunnersLag(ruleRunners);
                while (lag > 0) {
                    lag = getMaxLinkedRunnersLag(ruleRunners);
                    log.debug("Chain runner max lag={}, name={}", lag, chainName);
                }
            }, executor)
            .orTimeout(DEFAULT_MAX_LINKED_RUNNERS_CLOSE_WAIT_TIME_IN_SECONDS, TimeUnit.SECONDS)
            .exceptionally(e -> {
                log.error("Close linked runners timeout during stop chain runner, name={}", chainName, e);
                return null;
            })
            .thenRun(dispatchRunner::close)
            .exceptionally(e -> {
                log.error("Close dispatch runner failed, name={}", chainName, e);
                return null;
            });
    }

    public static List<PubSubRuleRunner> getLinkedPubSubRunners(Map<String, RuleRunner> ruleUidRunnerMap) {
        return ruleUidRunnerMap.values().stream()
            .map(e -> (PubSubRuleRunner) e)
            .collect(Collectors.toList());
    }

    public static int getMaxLinkedRunnersLag(List<PubSubRuleRunner> ruleRunners) {
        return ruleRunners.stream()
            .map(e -> e.getSubPendingSize())
            .max(Comparator.comparing(Integer::valueOf))
            .get();
    }
}
